import os
import re
from typing import Dict, List, Any
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks


# 全局变量
USER_NAME = 'A'
ROOT_DIR = f'C:/Users/{USER_NAME}/.cache/modelscope/hub'

# 命名实体识别模型来源：https://www.modelscope.cn/models/iic/nlp_raner_named-entity-recognition_chinese-base-cmeee/summary
MODEL = 'damo/nlp_raner_named-entity-recognition_chinese-base-cmeee'


# 模型路径
if os.path.exists(os.path.join(ROOT_DIR, MODEL)):
    model_dir = os.path.join(ROOT_DIR, MODEL)
else:
    model_dir = MODEL


# 模型加载
recognizer = pipeline(
    task=Tasks.named_entity_recognition, 
    model=model_dir, 
    show_progress_bar=True
)


# 用户输入（非常极端的例子）
query = '最近通常有胀痛、发炎，严重可出现【咽]\n喉炎、吞咽|{困难}，血$T3、T4升&高，TSH下降。' \
        '这种情况应该考虑可能是尘+螨过敏性//哮喘或者___代谢==性碱中毒，也可能是亚急性结---*（节性游走———）性脂膜炎，\n' \
        '建议前往皮--肤科或是心内**科就诊，在专业#@医师指导下做进一步检查。\n' \
        '最有可能采取 "药物治疗 康复治疗 及时补液，积极抗感染" 等措施。'

# 模型推理
result = recognizer(query)


# 去特殊字符
def remove_spedial_chars(text: str) -> str:
    special_chars = r'[\n\t\\(){}<>\[\]/|（）【】{}「」『』《》〈〉——“”‘’$#&@+*_=-]'
    return re.sub(special_chars, '', text)


# 提取知识图谱所需的实体类型、名称
def get_entities(result: Dict[str, List[Dict[str, Any]]]) -> List[Dict[str, str]]:
    entities = []
    for item in result.get('output'):
        if item['type'] != 'pro':  # 过滤pro实体
            item['span'] = remove_spedial_chars(item['span'])  # 去特殊字符
            entities.append({
                'type': item['type'],
                'span': item['span']
            })

    if len(entities) == 0:
        print('>>> 当前用户的输入内容中，没有您需要的实体!')
        return None
    
    return entities


all_entities = get_entities(result)
print(f'>>> 知识图谱所需实体信息:\n{all_entities}\n')  # 效果很不错